PREAMBLE = '''
INLINE size_t memtran(void *dest, void *src, size_t size) {
  memcpy(dest, src, size);
  return size;
}
#define pack_field(ptr, field, n) (ptr) += memtran((ptr), (field), sizeof(*(field)) * (n))
#define unpack_field(ptr, field, n) (ptr) += memtran((field), (ptr), sizeof(*(field)) * (n))
static int stag_next = 0x3001;
static int stag_prev = 0x3000;
static int rtag_next = 0x3000;
static int rtag_prev = 0x3001;
#define COMM_ISEND(mpp, size, dir, axis) \\
  MPI_Isend((mpp)->send_##dir, (size), MPI_BYTE, (mpp)->dir.axis, stag_##dir, (mpp)->comm, &(mpp)->send_req_##dir)
#define COMM_IRECV(mpp, dir, axis) \\
  MPI_Irecv((mpp)->recv_##dir, (mpp)->max_comm_size, MPI_BYTE, (mpp)->dir.axis, rtag_##dir, (mpp)->comm, &(mpp)->recv_req_##dir)
#define COMM_WAITALL(mpp)                                        \\
  {                                                              \\
    MPI_Wait(&((mpp)->send_req_prev), &((mpp)->send_stat_prev)); \\
    MPI_Wait(&((mpp)->send_req_next), &((mpp)->send_stat_next)); \\
    MPI_Wait(&((mpp)->recv_req_prev), &((mpp)->recv_stat_prev)); \\
    MPI_Wait(&((mpp)->recv_req_next), &((mpp)->recv_stat_next)); \\
  }
#define UP_ALIGN_PTR(ptr, alignment) ((void*)((((long)(ptr)) + (alignment - 1)) & ~(alignment - 1)))
#define UP_ALIGN_SIZE(size, alignment) (((size) + (alignment) - 1)) & ~ ((alignment) - 1))
#define MAX_CELL_COMM 49152
'''
print(PREAMBLE)
FIELD_TYPES= {
  'bonded_tag': "T*",
  'chain2_tag': "T*",
  'excl_tag': "T*",
  'scal_tag': "T*",
  'first_bonded': "i",
  'first_chain2': "i",
  'first_excl_atom': "i",
  'first_impr': "i",
  'first_scal_atom': "i",
  't': "I",
  'tag': "L",
  'impr_idx': "I*",
  'mass': "R",
  'rmass': "R",
  'q': "R",
  'f': "RVEC+",
  'v': "RVEC+",
  'x': "RVEC",
  'shake': "SHAKE",
  'shake_tmp': "RVEC+",
  'natom': "N",
  'nbonded_export': "N",
  'nchain2_export': "N",
  'nexcl_export': "N",
  'nexport': "N",
  'nimpr_export': "N",
  'nscal_export': "N",
}
TYPE_ALIGNMENT = {
  "T*": 8,
  "I": 4,
  'i': 4,
  "L": 8,
  "I*": 4,
  "R": 8,
  "RVEC+": 8,
  "RVEC": 8,
  "SHAKE": 8,
  "N": 4
}
TYPE_SIZE = {
  "T*": "sizeof(long)",
  "I": "sizeof(int)",
  'i': "sizeof(int)",
  "L": "sizeof(long)",
  "I*": "(sizeof(int) * 3)",
  "R": "sizeof(real)",
  "RVEC+": "(sizeof(real) * 3)",
  "RVEC": "(sizeof(real) * 3)",
  "SHAKE": "sizeof(shake_t)",
  "N": "sizeof(int)"
}
FIELD_INDEX = {
  "impr_idx": "first_impr",
  'bonded_tag': 'first_bonded',
  'chain2_tag': 'first_chain2',
  'scal_tag': 'first_scal_atom',
  'excl_tag': 'first_excl_atom'
}
def gen_estimate_size(name, fields, isexport = False):
  lines = [
    'size_t estimate_cell_size_%s(celldata_t *cell) {' % name,
    '  size_t ret = 0;'
  ]
  if isexport:
    lines.append('  int natoms = cell->nexport;')
    lines.append('  int atomst = CELL_CAP - cell->nexport;')
  else:
    lines.append('  int natoms = cell->natom;')
    lines.append('  int atomst = 0;')
  for field in fields:
    lines.append('  ret = UP_ALIGN_SIZE(ret, %d);' % TYPE_ALIGNMENT[FIELD_TYPES[field]])
    if FIELD_TYPES[field].endswith("*"):
      lines.append('  ret += (cell->%s[atomst + natoms] - cell->%s[atomst]) * %s;//%s' % (FIELD_INDEX[field], FIELD_INDEX[field], TYPE_SIZE[FIELD_TYPES[field]], field))
    elif FIELD_TYPES[field] == 'N':
      lines.append('  ret += 4; //%s' % field)
    elif FIELD_TYPES[field] == 'i':
      lines.append('  ret += (natoms + 1) * %s; //%s' % (TYPE_SIZE[FIELD_TYPES[field]], field))
    else:
      lines.append('  ret += natoms * %s;// %s' % (TYPE_SIZE[FIELD_TYPES[field]], field))
  lines.append('  return ret;')
  lines.append('}')
  return "\n".join(lines)

def gen_pack_cell_cpe(name, fields, isexport = False):
  lines = ['size_t pack_cell_%s_cpe(void *buf, celldata_t *cell){' % name]
  lines.extend([
    "  dma_init();"
    "  cellmeta_t meta;",
    "  pe_get(&cell->basis, &meta, sizeof(cellmeta_t));",
    "  dma_syn();",
    "  void lbuf[MAX_CELL_COMM];",
    "  void *lptr = lbuf + 8;",
  ])
  if isexport:
    lines.append('  int natoms = meta.nexport;')
    lines.append('  int atomst = CELL_CAP - meta.nexport;')
  else:
    lines.append('  int natoms = meta.natom;')
    lines.append('  int atomst = 0;')
  for field in fields:
    lines.append("  lptr = UP_ALIGN(lptr, %d);" % TYPE_ALIGNMENT[FIELD_TYPES[field]])
    if FIELD_TYPES[field] == "N":
      lines.append("  *(int*)lptr = meta.%s;" % field)
      lines.append("  lptr += 4;")
    elif FIELD_TYPES[field].endswith('*'):
        lines.append("  pe_get(cell->%s + (*%s_st), sizeof((*cell->%s)) * ((*%s_ed) - (*%s_st)));" % (field, FIELD_INDEX[field], field, FIELD_INDEX[field], FIELD_INDEX[field]))
        lines.append("  lptr += sizeof((*cell->%s)) * ((*%s_ed) - (*%s_st));" % (field, FIELD_INDEX[field], FIELD_INDEX[field]))
        lines.append("  dma_syn();")
    elif FIELD_TYPES[field] == 'i':
      lines.append("  int *%s_st = lptr;" % field)
      lines.append("  pe_get(cell->%s + atomst, sizeof((*cell->%s)) * (natoms + 1));" % (field, field))
      lines.append("  lptr += sizeof((*cell->%s)) * (natoms + 1);" % (field))
      lines.append("  int *%s_ed = lptr - 4;" % field)
      lines.append("  dma_syn();")
    else:
      lines.append("  pe_get(cell->%s + atomst, sizeof((*cell->%s)) * natoms);" % (field, field))
      lines.append("  lptr += sizeof((*cell->%s)) * natoms;" % (field))
      lines.append("  dma_syn();")
  lines.append('  lptr = UP_ALIGN(lptr, 32);')
  lines.append('  *(long*)lbuf = lptr - lbuf;')
  lines.append('  pe_put(buf, lbuf, lptr - buf);')
  lines.append('  dma_syn();')
  lines.append('  return lptr - lbuf;')
  lines.append('}')
  return "\n".join(lines)
def gen_unpack_cell_cpe(name, fields, isexport = False):
  lines = ['size_t unpack_cell_%s_cpe(void *buf, celldata_t *cell){' % name]
  lines.extend([
    "  dma_init();"
    "  cellmeta_t meta;",
    "  void lbuf[MAX_CELL_COMM];",
    "  pe_get(&cell->basis, &meta, sizeof(cellmeta_t));",
    "  pe_get(buf, lbuf, 256);",
    "  dma_syn();",
    "  size_t nbytes = *(long*)lbuf;",
    "  if (nbytes > 256) {",
    "  "
    "  }",
    "  void *lptr = lbuf;",
  ])
  if isexport:
    lines.append('  int natoms = meta.nexport;')
    lines.append('  int atomst = CELL_CAP - meta.nexport;')
  else:
    lines.append('  int natoms = meta.natom;')
    lines.append('  int atomst = 0;')
  for field in fields:
    lines.append("  lptr = UP_ALIGN(lptr, %d);" % TYPE_ALIGNMENT[FIELD_TYPES[field]])
    if FIELD_TYPES[field] == "N":
      lines.append("  *(int*)lptr = meta.%s;" % field)
      lines.append("  lptr += 4;")
    elif FIELD_TYPES[field].endswith('*'):
        lines.append("  pe_get(cell->%s + (*%s_st), sizeof((*cell->%s)) * ((*%s_ed) - (*%s_st)));" % (field, FIELD_INDEX[field], field, FIELD_INDEX[field], FIELD_INDEX[field]))
        lines.append("  lptr += sizeof((*cell->%s)) * ((*%s_ed) - (*%s_st));" % (field, FIELD_INDEX[field], FIELD_INDEX[field]))
        lines.append("  dma_syn();")
    elif FIELD_TYPES[field] == 'i':
      lines.append("  int *%s_st = lptr;" % field)
      lines.append("  pe_get(cell->%s + atomst, sizeof((*cell->%s)) * (natoms + 1));" % (field, field))
      lines.append("  lptr += sizeof((*cell->%s)) * (natoms + 1);" % (field))
      lines.append("  int *%s_ed = lptr - 4;" % field)
      lines.append("  dma_syn();")
    else:
      lines.append("  pe_get(cell->%s + atomst, sizeof((*cell->%s)) * natoms);" % (field, field))
      lines.append("  lptr += sizeof((*cell->%s)) * natoms;" % (field))
      lines.append("  dma_syn();")
  lines.append('  lptr = UP_ALIGN(lptr, 32);')
  lines.append('  pe_put(buf, lbuf, lptr - buf);')
  lines.append('  dma_syn();')
  lines.append('  return lptr - lbuf;')
  lines.append('}')
  return "\n".join(lines)
# def gen_pack_cell(name, fields, isexport = False):
#   lines = ['size_t pack_cell_%s(void *buf, celldata_t *cell){' % name]

#   if fields[0] in ['natom', 'nexport']:
#     lines.append('  *(long*)buf = cell->%s;' % fields[0])
#     lines.append('  void *ptr = buf + 8;')
#     fields = fields[1:]
#   else:
#     lines.append('  void *ptr = buf;')
#   for field in fields:
#     if field.startswith('n') and field.endswith('_export'):
#       lines.append('  *(long*)ptr = cell->%s;' % field)
#       lines.append('  ptr += 8;')
#     elif not isexport:
#       lines.append('  pack_field(ptr, cell->%s, cell->natom);' % field)
#     else:
#       if field in ['bonded_tag', 'chain2_tag', 'excl_tag', 'scal_tag', 'impr_idx']:
#         entry = field.split('_')[0]
#         lines.append('  pack_field(ptr, cell->%s + (MAX_%s_CELL - cell->n%s_export), cell->n%s_export);' % (field, entry.upper(), entry, entry))
#       elif field.startswith('first'):
#         lines.append('  pack_field(ptr, cell->%s + (CELL_CAP - cell->nexport), cell->nexport + 1);' % field)
#       else:
#         lines.append('  pack_field(ptr, cell->%s + (CELL_CAP - cell->nexport), cell->nexport);' % field)
#   lines.append('  return ptr - buf;')
#   lines.append('}')
#   return "\n".join(lines)

def gen_unpack_cell(name, fields, isexport = False):
  lines = ['size_t unpack_cell_%s(void *buf, celldata_t *cell){' % name]
  if fields[0] in ['natom', 'nexport']:
    lines.append('  cell->%s = *(long*)buf;' % fields[0])
    lines.append('  void *ptr = buf + 8;')
    fields = fields[1:]
  else:
    lines.append('  void *ptr = buf;')
  print('//' + str(fields))
  for field in fields:
    if field.startswith('n') and field.endswith('_export'):
      lines.append('  cell->%s = *(long*)ptr;' % field)
      lines.append('  ptr += 8;')
    elif field in ['f', 'v', 'shake_tmp']:
      if isexport:
        lines.append('  unpack_field(ptr, cell->%s + (CELL_CAP - cell->nexport), cell->nexport);' % field)
      elif name.startswith('reverse'):
        lines.append('  vec<real> *%s_buf = ptr;' % field)
        lines.append('  for (int i = 0; i < cell->natom; i ++){')
        lines.append('    vecaddv(cell->%s[i], cell->%s[i], %s_buf[i]);' % (field, field, field))
        lines.append('  }')
        lines.append('  ptr += sizeof(*(cell->%s)) * cell->natom;' % field)
      else:
        lines.append('  unpack_field(ptr, cell->%s, cell->natom);' % field)
    else:
      if isexport:
        if field in ['bonded_tag', 'chain2_tag', 'excl_tag', 'scal_tag', 'impr_idx']:
          entry = field.split('_')[0]
          lines.append('  unpack_field(ptr, cell->%s + (MAX_%s_CELL - cell->n%s_export), cell->n%s_export);' % (field, entry.upper(), entry, entry))
        elif field.startswith('first'):
          lines.append('  unpack_field(ptr, cell->%s + (CELL_CAP - cell->nexport), cell->nexport + 1);' % field)
        else:
          lines.append('  unpack_field(ptr, cell->%s + (CELL_CAP - cell->nexport), cell->nexport);' % field)
      else:
        lines.append('  unpack_field(ptr, cell->%s, cell->natom);' % field)
    
  lines.append('  return ptr - buf;')
  lines.append('}')
  return "\n".join(lines)

def gen_pack_unpack_brick(base):
  return """
size_t pack_brick_%s(void *buf, cellgrid_t *grid, int xlo, int xhi, int ylo, int yhi, int zlo, int zhi) {
  void *ptr = buf;
  for (int i = xlo; i < xhi; i ++) {
    for (int j = ylo; j < yhi; j ++) {
      for (int k = zlo; k < zhi; k ++) {
        celldata_t *cell = get_cell_xyz(grid, i, j, k);
        ptr += pack_cell_%s(ptr, cell);
      }
    }
  }
  return ptr - buf;
}

size_t unpack_brick_%s(void *buf, cellgrid_t *grid, int xlo, int xhi, int ylo, int yhi, int zlo, int zhi) {
  void *ptr = buf;
  for (int i = xlo; i < xhi; i ++) {
    for (int j = ylo; j < yhi; j ++) {
      for (int k = zlo; k < zhi; k ++) {
        celldata_t *cell = get_cell_xyz(grid, i, j, k);
        ptr += unpack_cell_%s(ptr, cell);
      }
    }
  }
  return ptr - buf;
}""" % (base, base, base, base)
def gen_pack_unpack_brick_cpe(base):
  return """
typedef struct pack_brick_param {
  void *buf;
  cellgrid_t *grid;
  int xlo, xhi, ylo, yhi, zlo, zhi;
} pack_brick_param_t;
#ifdef __sw_host__
extern void slave_pack_brick_%s_cpe(void *);
extern void slave_unpack_brick_%s_cpe(void *);
#endif
#ifdef __sw_slave__
#include "dma_macros_new.h"
#include <slave.h>
size_t pack_brick_%s_cpe(pack_brick_param_t *pm) {
  pack_brick_param_t lpm;
  cellgrid_t lgrid;
  pe_get(pm, &lpm, sizeof(pack_brick_param_t));
  dma_syn();
  pe_get(lpm.grid, &lgrid, sizeof(cellgrid_t));
  void *ptr = lpm.buf;
  int xlen = lpm.xhi - lpm.xlo;
  int ylen = lpm.yhi - lpm.ylo;
  int zlen = lpm.zhi - lpm.zlo;
  dma_syn();
  for (int i = xlo; i < xhi; i ++) {
    for (int j = ylo; j < yhi; j ++) {
      for (int k = zlo; k < zhi; k ++) {
        int cellid = ((i - xlo) * ylen + (j - ylo)) * zlen + k - zlo;
        if (cellid & 63 != _MYID) continue;
        celldata_t *cell = get_cell_xyz(&lgrid, i, j, k);
        ptr += pack_cell_%s(ptr, cell);
      }
    }
  }
  return ptr - buf;
}
size_t unpack_brick_%s_cpe(pack_brick_param_t *pm) {
  pack_brick_param_t lpm;
  cellgrid_t lgrid;
  pe_get(pm, &lpm, sizeof(pack_brick_param_t));
  dma_syn();
  pe_get(lpm.grid, &lgrid, sizeof(cellgrid_t));
  void *ptr = lpm.buf;
  int xlen = lpm.xhi - lpm.xlo;
  int ylen = lpm.yhi - lpm.ylo;
  int zlen = lpm.zhi - lpm.zlo;
  dma_syn();
  for (int i = xlo; i < xhi; i ++) {
    for (int j = ylo; j < yhi; j ++) {
      for (int k = zlo; k < zhi; k ++) {
        int cellid = ((i - xlo) * ylen + (j - ylo)) * zlen + k - zlo;
        if (cellid & 63 != _MYID) continue;
        celldata_t *cell = get_cell_xyz(&lgrid, i, j, k);
        ptr += unpack_cell_%s(ptr, cell);
      }
    }
  }
  return ptr - buf;
}

#endif""" % (base, base, base, base, base, base)

def gen_forward_comm(name):
  return """
void forward_comm_NAME(cellgrid_t *grid, mpp_t *mpp) {
  vec<int> *lo = &(grid->dim.lo);
  vec<int> *hi = &(grid->dim.hi);
  vec<int> *nlocal = &(grid->nlocal);
  int nn = grid->nn;

  size_t nsend_prev, nsend_next;

  COMM_IRECV(mpp, prev, z);
  COMM_IRECV(mpp, next, z);
  nsend_prev = pack_brick_forward_NAME(mpp->send_prev, grid, 0, nlocal->x, 0, nlocal->y, 0, nn);
  nsend_next = pack_brick_forward_NAME(mpp->send_next, grid, 0, nlocal->x, 0, nlocal->y, nlocal->z - nn, nlocal->z);
  COMM_ISEND(mpp, nsend_prev, prev, z);
  COMM_ISEND(mpp, nsend_next, next, z);
  COMM_WAITALL(mpp);
  unpack_brick_forward_NAME(mpp->recv_prev, grid, 0, nlocal->x, 0, nlocal->y, lo->z, lo->z + nn);
  unpack_brick_forward_NAME(mpp->recv_next, grid, 0, nlocal->x, 0, nlocal->y, hi->z - nn, hi->z);

  COMM_IRECV(mpp, prev, y);
  COMM_IRECV(mpp, next, y);
  nsend_prev = pack_brick_forward_NAME(mpp->send_prev, grid, 0, nlocal->x, 0, nn, lo->z, hi->z);
  nsend_next = pack_brick_forward_NAME(mpp->send_next, grid, 0, nlocal->x, nlocal->y - nn, nlocal->y, lo->z, hi->z);
  COMM_ISEND(mpp, nsend_prev, prev, y);
  COMM_ISEND(mpp, nsend_next, next, y);
  COMM_WAITALL(mpp);
  unpack_brick_forward_NAME(mpp->recv_prev, grid, 0, nlocal->x, lo->y, lo->y + nn, lo->z, hi->z);
  unpack_brick_forward_NAME(mpp->recv_next, grid, 0, nlocal->x, hi->y - nn, hi->y, lo->z, hi->z);

  COMM_IRECV(mpp, prev, x);
  COMM_IRECV(mpp, next, x);
  nsend_prev = pack_brick_forward_NAME(mpp->send_prev, grid, 0, nn, lo->y, hi->y, lo->z, hi->z);
  nsend_next = pack_brick_forward_NAME(mpp->send_next, grid, nlocal->x - nn, nlocal->x, lo->y, hi->y, lo->z, hi->z);
  COMM_ISEND(mpp, nsend_prev, prev, x);
  COMM_ISEND(mpp, nsend_next, next, x);
  COMM_WAITALL(mpp);
  unpack_brick_forward_NAME(mpp->recv_prev, grid, lo->x, lo->x + nn, lo->y, hi->y, lo->z, hi->z);
  unpack_brick_forward_NAME(mpp->recv_next, grid, hi->x - nn, hi->x, lo->y, hi->y, lo->z, hi->z);
}
""".replace('NAME', name)

def gen_reverse_comm(name):
  return """
void reverse_comm_NAME(cellgrid_t *grid, mpp_t *mpp) {
  vec<int> *lo = &(grid->dim.lo);
  vec<int> *hi = &(grid->dim.hi);
  vec<int> *nlocal = &(grid->nlocal);
  int nn = grid->nn;

  size_t nsend_prev, nsend_next;

  COMM_IRECV(mpp, prev, x);
  COMM_IRECV(mpp, next, x);
  nsend_prev = pack_brick_reverse_NAME(mpp->send_prev, grid, lo->x, lo->x + nn, lo->y, hi->y, lo->z, hi->z);
  nsend_next = pack_brick_reverse_NAME(mpp->send_next, grid, hi->x - nn, hi->x, lo->y, hi->y, lo->z, hi->z);
  COMM_ISEND(mpp, nsend_prev, prev, x);
  COMM_ISEND(mpp, nsend_next, next, x);
  COMM_WAITALL(mpp);
  unpack_brick_reverse_NAME(mpp->recv_prev, grid, 0, nn, lo->y, hi->y, lo->z, hi->z);
  unpack_brick_reverse_NAME(mpp->recv_next, grid, nlocal->x - nn, nlocal->x, lo->y, hi->y, lo->z, hi->z);

  COMM_IRECV(mpp, prev, y);
  COMM_IRECV(mpp, next, y);
  nsend_prev = pack_brick_reverse_NAME(mpp->send_prev, grid, 0, nlocal->x, lo->y, lo->y + nn, lo->z, hi->z);
  nsend_next = pack_brick_reverse_NAME(mpp->send_next, grid, 0, nlocal->x, hi->y - nn, hi->y, lo->z, hi->z);
  COMM_ISEND(mpp, nsend_prev, prev, y);
  COMM_ISEND(mpp, nsend_next, next, y);
  COMM_WAITALL(mpp);
  unpack_brick_reverse_NAME(mpp->recv_prev, grid, 0, nlocal->x, 0, nn, lo->z, hi->z);
  unpack_brick_reverse_NAME(mpp->recv_next, grid, 0, nlocal->x, nlocal->y - nn, nlocal->y, lo->z, hi->z);

  COMM_IRECV(mpp, prev, z);
  COMM_IRECV(mpp, next, z);
  nsend_prev = pack_brick_reverse_NAME(mpp->send_prev, grid, 0, nlocal->x, 0, nlocal->y, lo->z, lo->z + nn);
  nsend_next = pack_brick_reverse_NAME(mpp->send_next, grid, 0, nlocal->x, 0, nlocal->y, hi->z - nn, hi->z);
  COMM_ISEND(mpp, nsend_prev, prev, z);
  COMM_ISEND(mpp, nsend_next, next, z);
  COMM_WAITALL(mpp);
  unpack_brick_reverse_NAME(mpp->recv_prev, grid, 0, nlocal->x, 0, nlocal->y, 0, nn);
  unpack_brick_reverse_NAME(mpp->recv_next, grid, 0, nlocal->x, 0, nlocal->y, nlocal->z - nn, nlocal->z);
}
""".replace('NAME', name)

fields = {
  #'forward_bond': ['bonded_id'],
  'forward_most': ['natom', 'tag', 'x', 'q', 't', 'mass', 'rmass'],
  # 'forward_rebuild': ['natom', 'tag', 'x', 'q', 't', 'bonded_tag', 'excl_tag', 'scal_tag'],
  'forward_x': ['x'],
  'reverse_f': ['f'],
  'forward_v': ['v'],
  'reverse_v': ['v'],
  # 'forward_export': ['nexport', 'x', 'q', 'tag', 't', 'mass', 'bonded_tag', 'excl_tag', 'scal_tag', 'chain2_tag', 'impr_idx'],
  'forward_export_list': ['nexport', 'x', 'q', 'tag', 't', 'v', 'mass', # 'bonded_tag', 'excl_tag', 'scal_tag', 'chain2_tag', 'impr_idx',
    'nbonded_export', 'nchain2_export', 'nscal_export', 'nexcl_export', 'nimpr_export',
    'first_bonded', 'first_chain2', 'first_scal_atom', 'first_excl_atom', 'first_impr',
    'bonded_tag', 'chain2_tag', 'excl_tag', 'scal_tag', 'impr_idx', 'shake'],
  'forward_export_list_cg': ['nexport', 'x', 'q', 'tag', 't', 'v', 'mass', # 'bonded_tag', 'excl_tag', 'scal_tag', 'chain2_tag', 'impr_idx',
    'nbonded_export', 'nchain2_export', 'nscal_export', 'nexcl_export', 'nimpr_export',
    'first_bonded', 'first_chain2', 'first_scal_atom', 'first_excl_atom', 'first_impr',
    'bonded_tag', 'chain2_tag', 'excl_tag', 'scal_tag', 'impr_idx', 'shake', 'shake_tmp'],
  'forward_shake': ['shake_tmp'],
  'reverse_shake': ['shake_tmp']
}
for k, v in fields.items():
  print(gen_estimate_size(k, v, v[0] == 'nexport'))
  print(gen_pack_cell_cpe(k, v, v[0] == 'nexport'))
  # print(gen_pack_cell(k, v, v[0] == 'nexport'))
  # print(gen_unpack_cell(k, v, v[0] == 'nexport'))
  # print(gen_pack_unpack_brick(k))
  # print(gen_pack_unpack_brick_cpe(k))
  # if k.startswith('forward'):
  #   print(gen_forward_comm(k[8:]))
  # else:
  #   print(gen_reverse_comm(k[8:]))
